/*
* Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

use proc_macro::TokenStream;
use syn::{Data, DeriveInput, GenericArgument, Lifetime, PathArguments};

#[proc_macro_derive(Encode, attributes(encode))]
pub fn derive_encode(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as DeriveInput);
    match &input.data {
        Data::Struct(data_struct) => {
            let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
            let ty = &input.ident;

            let encode_unchecked_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            unsafe { self.#ident.encode_unchecked(encoder, buf) };
                        })
                    } else {
                        None
                    }
                });
            let encode_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            self.#ident.encode(encoder, buf)?;
                        })
                    } else {
                        None
                    }
                });

            quote::quote! {
                impl #impl_generics ohmw_core::bincode::Encode for #ty #ty_generics #where_clause {
                    unsafe fn encode_unchecked<E: ohmw_core::bincode::Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
                        #(#encode_unchecked_stmts)*
                    }

                    fn encode<E: ohmw_core::bincode::Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
                        #(#encode_stmts)*
                        Ok(())
                    }
                }
            }.into()
        },
        _ => panic!("Only support struct"),
    }
}

#[proc_macro_derive(Decode, attributes(decode))]
pub fn derive_decode(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as DeriveInput);
    match &input.data {
        Data::Struct(data_struct) => {
            let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
            let ty = &input.ident;

            let decode_unchecked_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        let ty = &field.ty;
                        Some(quote::quote! {
                            let #ident = unsafe { <#ty as ohmw_core::bincode::Decode>::decode_unchecked(decoder, buf) };
                        })
                    } else {
                        None
                    }
                });
            let decode_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        let ty = &field.ty;
                        Some(quote::quote! {
                            let #ident = <#ty as ohmw_core::bincode::Decode>::decode(decoder, buf)?;
                        })
                    } else {
                        None
                    }
                });
            let fields_with_comma = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            #ident,
                        })
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>();

            quote::quote! {
                impl #impl_generics ohmw_core::bincode::Decode for #ty #ty_generics #where_clause {
                    unsafe fn decode_unchecked<D: ohmw_core::bincode::Decoder>(decoder: &mut D, buf: &[u8]) -> Self {
                        #(#decode_unchecked_stmts)*
                        Self {
                            #(#fields_with_comma)*
                        }
                    }

                    fn decode<D: ohmw_core::bincode::Decoder>(decoder: &mut D, buf: &[u8]) -> Result<Self, D::Error> {
                        #(#decode_stmts)*
                        Ok(Self {
                            #(#fields_with_comma)*
                        })
                    }
                }
            }.into()
        },
        _ => panic!("Only support struct"),
    }
}

#[proc_macro_derive(BorrowDecode, attributes(borrow_decode))]
pub fn derive_borrow_decode(input: TokenStream) -> TokenStream {
    let mut input = syn::parse_macro_input!(input as DeriveInput);
    match &mut input.data {
        Data::Struct(data_struct) => {
            let ty = &input.ident;
            input.generics.lifetimes_mut().for_each(|l| {
                l.lifetime = Lifetime::new("'de", proc_macro2::Span::call_site());
            });
            data_struct.fields.iter_mut().for_each(|f| {
                rename_lifetime(&mut f.ty, "'de");
            });
            let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

            let borrow_decode_unchecked_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        let ty = &field.ty;
                        Some(quote::quote! {
                            let #ident = unsafe { <#ty as ohmw_core::bincode::BorrowDecode<'de>>::borrow_decode_unchecked(decoder, buf) };
                        })
                    } else {
                        None
                    }
                });
            let borrow_decode_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        let ty = &field.ty;
                        Some(quote::quote! {
                            let #ident = <#ty as ohmw_core::bincode::BorrowDecode<'de>>::borrow_decode(decoder, buf)?;
                        })
                    } else {
                        None
                    }
                });
            let fields_with_comma = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            #ident,
                        })
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>();

            quote::quote! {
                impl #impl_generics ohmw_core::bincode::BorrowDecode<'de> for #ty #ty_generics #where_clause {
                    unsafe fn borrow_decode_unchecked<D: ohmw_core::bincode::Decoder>(decoder: &mut D, buf: &'de [u8]) -> Self {
                        #(#borrow_decode_unchecked_stmts)*
                        Self {
                            #(#fields_with_comma)*
                        }
                    }

                    fn borrow_decode<D: ohmw_core::bincode::Decoder>(decoder: &mut D, buf: &'de [u8]) -> Result<Self, D::Error> {
                        #(#borrow_decode_stmts)*
                        Ok(Self {
                            #(#fields_with_comma)*
                        })
                    }
                }
            }.into()
        },
        _ => panic!("Only support struct"),
    }
}

fn rename_lifetime(ty: &mut syn::Type, symbol: &str) {
    match ty {
        syn::Type::Array(array) => {
            rename_lifetime(&mut array.elem, symbol);
        },
        syn::Type::Paren(paren) => {
            rename_lifetime(&mut paren.elem, symbol);
        },
        syn::Type::Reference(reference) => {
            if let Some(lifetime) = reference.lifetime.as_mut() {
                *lifetime = Lifetime::new(symbol, proc_macro2::Span::call_site());
            }
        },
        syn::Type::Slice(slice) => {
            rename_lifetime(&mut slice.elem, symbol);
        },
        syn::Type::Tuple(tuple) => {
            tuple.elems.iter_mut().for_each(|e| rename_lifetime(e, symbol));
        },
        syn::Type::Path(path) => {
            if let Some(qself) = path.qself.as_mut() {
                rename_lifetime(&mut qself.ty, symbol);
            }
            path.path.segments.iter_mut().for_each(|s| {
                match &mut s.arguments {
                    PathArguments::AngleBracketed(bracketed) => {
                        bracketed.args.iter_mut().for_each(|a| {
                            match a {
                                GenericArgument::Type(ty) => {
                                    rename_lifetime(ty, symbol);
                                },
                                GenericArgument::Lifetime(lifetime) => {
                                    *lifetime = Lifetime::new(symbol, proc_macro2::Span::call_site());
                                },
                                _ => {}
                            }
                        });
                    },
                    _ => {}
                }
            });
        },
        _ => {}
    }
}
